package com.example.springboot.config.aop;

import com.example.springboot.config.anotation.requstlimit.Limit;
import com.google.common.collect.Maps;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Map;

/**
 * @Author:ChenZhangKun
 * @Date: 2021/10/11 15:13
 */
@Component
@Aspect
@Slf4j
public class ControllerLimitAop {
    @Pointcut(value = "execution(* com.example.springboot.controller..*Controller.*(..))")
    public void aopPoint() {

    }

    /**
     * 不同的接口，不同的流量控制
     * map的key为 Limiter.key
     */
    private final Map<String, RateLimiter> limitMap = Maps.newConcurrentMap();

    @Around("aopPoint()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        log.info("【进入guava令牌桶算法的aop切面了】");
        // 拿到方法签名
        MethodSignature methodSignature = (MethodSignature) joinPoint.getSignature();
        Method method = methodSignature.getMethod();
        // 是否被limit注解
        if (!method.isAnnotationPresent(Limit.class)) {
            log.info("【没有接口限流，放行】");
            return joinPoint.proceed();
        }
        // 拿到注解
        Limit limit = method.getAnnotation(Limit.class);
        // 那大炮参数
        String key = limit.key();
        if (!limitMap.containsKey(key)) {
            // 创建限流措施
            double permitsPerSecond = limit.permitsPerSecond();
            RateLimiter rateLimiter = RateLimiter.create(permitsPerSecond);
            // 放入
            limitMap.put(key, rateLimiter);
            String interfaceName = getInterfaceName();
            log.info("【新建接口:{}的令牌桶】", interfaceName);
        }
        RateLimiter rateLimiter = limitMap.get(key);
        // 拿到令牌
        boolean tryAcquire = rateLimiter.tryAcquire(limit.timeout(), limit.timeunit());
        if (!tryAcquire) {
            log.info("令牌桶={}，获取令牌失败", key);
            return limit.msg();
        }
        log.info("【获取接口：{}】令牌桶成功", getInterfaceName());
        return joinPoint.proceed();
    }

    private String getInterfaceName() {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        HttpServletRequest httpServletRequest = (HttpServletRequest) requestAttributes.resolveReference(RequestAttributes.REFERENCE_REQUEST);
        return httpServletRequest.getRequestURI();
    }
}
